# Copyright (c) Huawei Technologies Co., Ltd. 2023. All rights reserved.
from typing import Union, List
import random
from atk.case_generator.generator.generate_types import GENERATOR_REGISTRY
from atk.case_generator.generator.base_generator import CaseGenerator
from atk.configs.case_config import InputCaseConfig, CaseConfig

@GENERATOR_REGISTRY.register("ascend_generate_aclnn_moe_token_unpermute_with_ep")
class DefaultGenerator(CaseGenerator):
    def __init__(self, config):
        super().__init__(config)
        self.tokens_num = 1
        self.top_k = 1
        self.hidden_size = 1
        self.total_len = 1
        self.start = 1
        self.end = 1

    def after_input_config(
            self,
            index: int,
            input_case: Union[InputCaseConfig, List[InputCaseConfig]],
            ) -> Union[InputCaseConfig, List[InputCaseConfig]]:
        if index == 0:
            if input_case.shape == []:
                input_case.shape = [7619, 7680]
            self.hidden_size = input_case.shape[0]
        if index == 1:
            if input_case.shape == []:
                input_case.shape = [16384]
            self.total_len = input_case.shape[0]
        if index == 2:
            if input_case.shape != []:
                self.top_k = input_case.shape[1]
                self.tokens_num = int(self.total_len / self.top_k)
            else:
                self.tokens_num = self.total_len
                self.top_k = 1
        if index == 4:
            self.start = random.randint(0, self.total_len)
            self.end = random.randint(0, self.total_len)
            if self.end < self.start:
                self.start, self.end = self.end, self.start
            if self.start == self.end:
                self.start = 0
                self.end = self.total_len
            input_case[0].range_values = self.start
            input_case[1].range_values = self.end

        return input_case

    def after_case_config(self, case_config: CaseConfig) -> CaseConfig:
        case_config.inputs[0].shape = [self.tokens_num, self.hidden_size]
        case_config.inputs[1].shape = [self.tokens_num * self.top_k]
        case_config.inputs[2].shape = [self.tokens_num, self.top_k]
        return case_config

